'''
conda activate py39
python main.py
'''
from retinaface_detector import RetinaFaceDetector
from sface_feature_extractor import sFaceFeatureExtractor
#from face_det_feature import FaceDetFeature
from db_operate import FaceDB

import cv2
import numpy as np
import json
from datetime import datetime, timezone
import time
import os
from copy import deepcopy


from yolox.tracker.byte_tracker import BYTETracker
class Args:
    def __init__(self):
        # 模型相关参数
        self.model = "yolox_s"
        self.ckpt = "yolox_s.pth"
        # 跟踪器参数
        self.track_thresh = 0.5
        self.track_buffer = 30
        self.match_thresh = 0.8
        self.min_box_area = 100
        self.mot20 = False
        
        # 其他参数
        self.device = "gpu"
        self.fp16 = False
        self.fuse = False
# 创建参数对象
args = Args() # 配置ByteTrack的参数
# 针对RK3588的优化配置
args = Args()
args.track_thresh = 0.6      # 提高阈值，减少低置信度检测框的处理
args.track_buffer = 20       # 减少轨迹保留帧数，降低内存消耗
args.match_thresh = 0.85     # 提高匹配阈值，减少错误匹配
args.min_box_area = 5      # 增大最小检测框面积，过滤小目标
args.device = "cpu"          # 如果GPU性能不足，使用CPU
args.fp16 = True             # 使用FP16精度，减少计算量
args.fuse = True             # 融合模型，提高推理速度
# 初始化ByteTrack跟踪器
tracker = BYTETracker(args)
# 模拟图像信息
img_info = [480, 640]  # 图像高度和宽度
img_size = [480, 640]

target_buf = {}

from numpy.linalg import norm
def compare_features(feature1, feature2):
    """计算余弦相似度判断是否为同一人"""
    cos_sim = np.dot(feature1, feature2) / (norm(feature1) * norm(feature2))
    print(cos_sim)
    return cos_sim

def main():
    faceDB = FaceDB("face.db")
    faceID = {}
    # 同步数据库所有人脸ID信息
    rows = faceDB.select_all_faceid()
    for row in rows:
        faceID[row[0]] = {"name":row[1], "num":row[2], "feature":np.array(json.loads(row[3]))}
    
    #faceDetFeature = FaceDetFeature()
    # 初始化人脸检测器
    face_detector = RetinaFaceDetector()  # target=None
    # 初始化特征提取器
    feature_extractor = sFaceFeatureExtractor()  #target=None
    
    cap = cv2.VideoCapture(21)
    # 检查摄像头是否成功打开
    if not cap.isOpened():
        print("无法打开摄像头")
        sys.exit()
    print("按 'q' 键退出")
    
    while True:
        # 获取时间戳
        timestamp_ms = datetime.now(timezone.utc).timestamp() # 精确到毫秒  毫秒用小数点表示
        timestamp = int(timestamp_ms*1000)
        # 先加上所有目标不存在了，检测发现了再标记存在
        for key in target_buf:
            target_buf[key]["state"] = 0
            
        # 读取一帧图像
        ret, frame = cap.read()
        # 检查是否成功读取图像
        if not ret:
            print("无法获取图像")
            break
        img = deepcopy(frame)
        #print(img.shape)
        #coords,features = faceDetFeature.get_info(img)
        coords = []
        features = []
        faces  = face_detector.detect(img, score_threshold=0.5, nms_threshold=0.5)
        if len(faces ) > 0:
            # 提取人脸坐标
            for face in faces:
                x1, y1, x2, y2 = map(int, face[:4])  # 取第一个人脸的坐标
                score = face[4]
                coords.append([x1, y1, x2, y2, score])
        #else:
        #    print("未检测到人脸")
        if len(coords) != 0:
            # 更新跟踪器
            coords_array = np.array(coords,dtype=float)
            dets = np.concatenate((coords_array[:,:4].astype(int), coords_array[:,4].reshape([-1,1])), axis=1)
            online_targets = tracker.update(np.array(coords,dtype=float), img_info, img_size)
            # 处理跟踪结果
            for target in online_targets:
                ltwh = target.tlwh.astype(int)
                tid = target.track_id
                score = target.score
                #print(f"Track ID: {tid}, Bbox: {ltwh}, Score: {score}")
                left = ltwh[0]
                top = ltwh[1]
                right = ltwh[0]+ltwh[2]
                bottom = ltwh[1]+ltwh[3]
                # 找到最合适人脸角度，只抓拍有一次
                if tid not in target_buf.keys():
                    target_buf[tid] = {"coord":[left, top, right, bottom], "score":score, "img":frame, "state": 1, "s_time":time.time()}
                else:
                    # 更大目标，匹配上
                    #if right+bottom>target_buf[tid]["coord"][2]+target_buf[tid]["coord"][3]:
                    if score>target_buf[tid]["score"]:
                        target_buf[tid] = {"coord":[left, top, right, bottom], 
                                           "score":score, 
                                           "img":frame, 
                                           "state": 1, 
                                           "s_time":target_buf[tid]["s_time"]}
                    else:
                        target_buf[tid]["state"] = 1
                cv2.rectangle(img, (left, top), (right, bottom), (0, 255, 0), 2)
                cv2.putText(img, f"Face: {tid}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        # 根据目标状态来确定是否入库
        # for key in target_buf:  这样不能删除，不能用。
        for key_target_buf in list(target_buf.keys()):
            if target_buf[key_target_buf]["state"] == 0:
                img_cur = target_buf[tid]["img"]
                coord = target_buf[tid]["coord"]
                f_array = feature_extractor.extract(img_cur, [coord])[0]
                f_list = [float(x) for x in f_array]  # list(f_array) np.float32不能直接转
                f_json = json.dumps(f_list) #列表才能被序列化
                c_list = [float(x) for x in coord] 
                c_json = json.dumps(c_list)
                # 找出ID，没有匹配就新增
                name = ""
                score = 0.0
                for key in faceID:
                    score_temp = compare_features(faceID[key]["feature"], f_array)
                    if score_temp>score:
                        score = score_temp
                        name = faceID[key]["name"]
                if score<0.4:
                    # 数据库里面没有对应人脸
                    cur_id = faceDB.get_faceid_max_id()+1
                    name = str(cur_id)
                    # 插入一条人脸ID
                    faceDB.insert_faceid((name, 1, f_json))
                    # 同步更新缓存
                    faceID[cur_id] = {"name":name, "num":0, "feature":f_array}
                    rows = faceDB.select_all_faceid()
                    print("所有ID:")
                    for row in rows:
                        print(row[0],row[1],row[2])
                    print("\n")
                # 记录人脸
                faceDB.insert_faceRecord((f'./save_img/{timestamp}.jpg', name, c_json, f_json))
                cv2.imwrite(f'./save_img/{timestamp}.jpg', img_cur)
                rows = faceDB.select_all_faceid()
                print("所有人脸ID:")
                for row in rows:
                    print(row[0],row[1],row[2])
                print("\n")
                rows = faceDB.select_all_faceRecord()
                print("所有人脸记录:")
                for row in rows:
                    print(row[0],row[1],row[2],np.array(json.loads(row[3])),np.array(json.loads(row[4])).shape)
                print("\n")
                # 删除消失人脸批号
                del target_buf[key_target_buf]
        #print(list(target_buf.keys()))
        # 测试
        timestamp = int(datetime.now(timezone.utc).timestamp()*1000) # 精确到毫秒
        for f_array, coord in zip(features, coords):
            x1, y1, x2, y2 = coord[:4]
            #cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            #cv2.putText(img, f"Face: {name}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        # 显示图像
        cv2.imshow('USB Camera', img)
        # 按'q'键退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    '''
    #image_path = "./test_img/test.jpg"  # 测试图像
    file_paths = []
    for entry in os.listdir("./test_img1"):
        entry_path = os.path.join("./test_img1", entry)
        if os.path.isfile(entry_path):
            file_paths.append(entry_path)
    
    for image_path in file_paths:
        img = cv2.imread(image_path)
        coords,features = faceDetFeature.get_info(img)
        
        # 测试
        timestamp = int(datetime.now(timezone.utc).timestamp()*1000) # 精确到毫秒
        for f_array, coord in zip(features, coords):
            f_list = [float(x) for x in f_array]  # list(f_array) np.float32不能直接转
            f_json = json.dumps(f_list) #列表才能被序列化
            c_json = json.dumps(coord)
            # 找出ID，没有匹配就新增
            name = ""
            score = 0.0
            for key in faceID:
                score_temp = faceDetFeature.compare_features(faceID[key]["feature"], f_array)
                if score_temp>score:
                    score = score_temp
                    name = faceID[key]["name"]
            if score<0.4:
                # 数据库里面没有对应人脸
                cur_id = faceDB.get_faceid_max_id()+1
                name = str(cur_id)
                # 插入一条人脸ID
                faceDB.insert_faceid((name, 1, f_json))
                # 同步更新缓存
                faceID[cur_id] = {"name":name, "num":0, "feature":f_array}
                rows = faceDB.select_all_faceid()
                print("所有ID:")
                for row in rows:
                    print(row[0],row[1],row[2])
                print("\n")
            # 记录人脸
            faceDB.insert_faceRecord((f'./save_img/{timestamp}.jpg', name, c_json, f_json))
        cv2.imwrite(f'./save_img/{timestamp}.jpg', img)
        rows = faceDB.select_all_faceRecord()
        print("所有记录:")
        for row in rows:
            print(row[0],row[1],row[2],np.array(json.loads(row[3])),np.array(json.loads(row[4])).shape)
        print("\n")
    '''




if __name__ == "__main__":
    main()


